# @Time : 2021/8/7 15:58
# @Author : Li Kunlun
# @Description : 动量法
import utils as d2l
from mxnet import nd

# 1、梯度下降的问题
eta = 0.4


# 构造一个输入为二维向量 𝒙=[𝑥1,𝑥2]和
# 输出为标量的目标函数 𝑓(𝒙)=0.1*(𝑥_1)^2+2*(x_2)^2,(𝑥_1)^2平方系数变小
def f_2d(x1, x2):
    return 0.1 * x1 ** 2 + 2 * x2 ** 2


def gd_2d(x1, x2, s1, s2):
    return (x1 - eta * 0.2 * x1, x2 - eta * 4 * x2, 0, 0)


# todo 目标函数在竖直方向（ 𝑥2 轴方向）比在水平方向（ 𝑥1 轴方向）的斜率的绝对值更大,
# 将学习率调大时，此时自变量在竖直方向不断越过最优解并逐渐发散
d2l.show_trace_2d(f_2d, d2l.train_2d(gd_2d))


# 2、动量法
# 对上面的梯度下降法，使用动量法进行改进
def momentum_2d(x1, x2, v1, v2):
    v1 = gamma * v1 + eta * 0.2 * x1
    v2 = gamma * v2 + eta * 4 * x2
    return x1 - v1, x2 - v2, v1, v2


eta, gamma = 0.4, 0.5
# 可以看到使用较小的学习率 𝜂=0.4 和动量超参数 𝛾=0.5 时，
# 动量法在竖直方向上的移动更加平滑，且在水平方向上更快逼近最优解。
d2l.show_trace_2d(f_2d, d2l.train_2d(momentum_2d))

# 从零开始实现
features, labels = d2l.get_data_ch7()


# 动量法需要对每一个自变量维护一个同它一样形状的速度变量，且超参数里多了动量超参数
def init_momentum_states():
    v_w = nd.zeros((features.shape[1], 1))
    v_b = nd.zeros(1)
    return (v_w, v_b)


def sgd_momentum(params, states, hyperparams):
    for p, v in zip(params, states):
        v[:] = hyperparams['momentum'] * v + hyperparams['lr'] * p.grad
        p[:] -= v


# 小批量随机梯度为最近2个时间步的2倍小批量梯度的加权平均。
d2l.train_ch7(sgd_momentum, init_momentum_states(), {'lr': 0.02, 'momentum': 0.5}, features, labels)

d2l.train_ch7(sgd_momentum, init_momentum_states(), {'lr': 0.02, 'momentum': 0.9}, features, labels)

# 将学习率减小到原来的1/5。目标函数值在下降了一段时间后变化更加平滑。
d2l.train_ch7(sgd_momentum, init_momentum_states(), {'lr': 0.004, 'momentum': 0.9}, features, labels)

print("--------------------简洁实现--------------------------")
# 在Gluon中，只需要在Trainer实例中通过momentum来指定动量超参数即可使用动量法
d2l.train_gluon_ch7('sgd', {'learning_rate': 0.004, 'momentum': 0.9}, features, labels)
